// Copyright Epic Games, Inc. All Rights Reserved.

#pragma once

#include <zencore/compactbinary.h>
#include <zencore/uid.h>
#include <zencore/xxhash.h>
#include <zenhttp/httpserver.h>
#include <zenstore/gc.h>

ZEN_THIRD_PARTY_INCLUDES_START
#include <tsl/robin_map.h>
ZEN_THIRD_PARTY_INCLUDES_END

#include <map>
#include <unordered_map>

namespace zen {

class CbPackage;
class CidStore;
class AuthMgr;
class ScrubContext;
class JobQueue;

enum class HttpResponseCode;

struct OplogEntry
{
	uint32_t OpLsn;
	uint32_t OpCoreOffset;	// note: Multiple of alignment!
	uint32_t OpCoreSize;
	uint32_t OpCoreHash;  // Used as checksum
	Oid		 OpKeyHash;
	uint32_t Reserved;

	inline bool IsTombstone() const { return OpCoreOffset == 0 && OpCoreSize == 0 && OpLsn == 0; }
	inline void MakeTombstone() { OpLsn = OpCoreOffset = OpCoreSize = OpCoreHash = Reserved = 0; }
};

struct OplogEntryAddress
{
	uint64_t Offset;
	uint64_t Size;
};

static_assert(IsPow2(sizeof(OplogEntry)));

/** Project Store

	A project store consists of a number of Projects.

	Each project contains a number of oplogs (short for "operation log"). UE uses
	one oplog per target platform to store the output of the cook process.

	An oplog consists of a sequence of "op" entries. Each entry is a structured object
	containing references to attachments. Attachments are typically the serialized
	package data split into separate chunks for bulk data, exports and header
	information.
 */
class ProjectStore : public RefCounted, public GcStorage, public GcContributor, public GcReferencer, public GcReferenceLocker
{
	struct OplogStorage;

public:
	ProjectStore(CidStore& Store, std::filesystem::path BasePath, GcManager& Gc, JobQueue& JobQueue);
	~ProjectStore();

	struct Project;

	struct Oplog
	{
		Oplog(std::string_view			   Id,
			  Project*					   Project,
			  CidStore&					   Store,
			  std::filesystem::path		   BasePath,
			  const std::filesystem::path& MarkerPath);
		~Oplog();

		[[nodiscard]] static bool ExistsAt(const std::filesystem::path& BasePath);

		void Read();
		void Write();
		void Update(const std::filesystem::path& MarkerPath);
		bool Reset();

		struct ChunkInfo
		{
			Oid		 ChunkId;
			uint64_t ChunkSize;
		};

		std::vector<ChunkInfo> GetAllChunksInfo();
		void				   IterateChunkMap(std::function<void(const Oid&, const IoHash& Hash)>&& Fn);
		void   IterateFileMap(std::function<void(const Oid&, const std::string_view& ServerPath, const std::string_view& ClientPath)>&& Fn);
		void   IterateOplog(std::function<void(CbObjectView)>&& Fn);
		void   IterateOplogWithKey(std::function<void(int, const Oid&, CbObjectView)>&& Fn);
		void   IterateOplogLocked(std::function<void(CbObjectView)>&& Fn);
		size_t GetOplogEntryCount() const;

		std::optional<CbObject> GetOpByKey(const Oid& Key);
		std::optional<CbObject> GetOpByIndex(int Index);
		int						GetOpIndexByKey(const Oid& Key);
		int						GetMaxOpIndex() const;

		IoBuffer					 FindChunk(const Oid& ChunkId);
		IoBuffer					 GetChunkByRawHash(const IoHash& RawHash);
		bool						 IterateChunks(std::span<IoHash>												 RawHashes,
												   const std::function<bool(size_t Index, const IoBuffer& Payload)>& AsyncCallback,
												   WorkerThreadPool*												 OptionalWorkerPool);
		bool						 IterateChunks(std::span<Oid>													 ChunkIds,
												   const std::function<bool(size_t Index, const IoBuffer& Payload)>& AsyncCallback,
												   WorkerThreadPool*												 OptionalWorkerPool);
		inline static const uint32_t kInvalidOp = ~0u;

		/** Persist a new oplog entry
		 *
		 * Returns the oplog LSN assigned to the new entry, or kInvalidOp if the entry is rejected
		 */
		uint32_t AppendNewOplogEntry(CbPackage Op);

		uint32_t			  AppendNewOplogEntry(CbObjectView Core);
		std::vector<uint32_t> AppendNewOplogEntries(std::span<CbObjectView> Cores);

		enum UpdateType
		{
			kUpdateNewEntry,
			kUpdateReplay
		};

		const std::string& OplogId() const { return m_OplogId; }

		const std::filesystem::path& TempPath() const { return m_TempPath; }
		const std::filesystem::path& MarkerPath() const { return m_MarkerPath; }

		LoggerRef		Log() { return m_OuterProject->Log(); }
		void			Flush();
		void			ScrubStorage(ScrubContext& Ctx);
		void			GatherReferences(GcContext& GcCtx);
		static uint64_t TotalSize(const std::filesystem::path& BasePath);
		uint64_t		TotalSize() const;

		std::size_t OplogCount() const
		{
			RwLock::SharedLockScope _(m_OplogLock);
			return m_LatestOpMap.size();
		}

		std::filesystem::path PrepareForDelete(bool MoveFolder);

		void AddChunkMappings(const std::unordered_map<Oid, IoHash, Oid::Hasher>& ChunkMappings);

		void CaptureUpdatedLSNs(std::span<const uint32_t> LSNs);
		void CaptureAddedAttachments(std::span<const IoHash> AttachmentHashes);

		void					EnableUpdateCapture();
		void					DisableUpdateCapture();
		void					IterateCapturedLSNs(std::function<bool(const CbObjectView& UpdateOp)>&& Callback);
		std::vector<IoHash>		GetCapturedAttachments();
		RwLock::SharedLockScope GetGcReferencerLock() { return RwLock::SharedLockScope(m_OplogLock); }

	private:
		struct FileMapEntry
		{
			std::string ServerPath;
			std::string ClientPath;
		};

		template<class V>
		using OidMap = tsl::robin_map<Oid, V, Oid::Hasher>;

		Project*			  m_OuterProject = nullptr;
		CidStore&			  m_CidStore;
		std::filesystem::path m_BasePath;
		std::filesystem::path m_MarkerPath;
		std::filesystem::path m_TempPath;

		mutable RwLock						   m_OplogLock;
		OidMap<IoHash>						   m_ChunkMap;		   // output data chunk id -> CAS address
		OidMap<IoHash>						   m_MetaMap;		   // meta chunk id -> CAS address
		OidMap<FileMapEntry>				   m_FileMap;		   // file id -> file map entry
		int32_t								   m_ManifestVersion;  // File system manifest version
		tsl::robin_map<int, OplogEntryAddress> m_OpAddressMap;	   // Index LSN -> op data in ops blob file
		OidMap<int>							   m_LatestOpMap;	   // op key -> latest op LSN for key

		mutable RwLock						   m_UpdateCaptureLock;
		uint32_t							   m_UpdateCaptureRefCounter = 0;
		std::unique_ptr<std::vector<uint32_t>> m_CapturedLSNs;
		std::unique_ptr<std::vector<IoHash>>   m_CapturedAttachments;

		RefPtr<OplogStorage> m_Storage;
		const std::string	 m_OplogId;

		RefPtr<OplogStorage> GetStorage();

		/** Scan oplog and register each entry, thus updating the in-memory tracking tables
		 */
		void ReplayLog();

		struct OplogEntryMapping
		{
			struct Mapping
			{
				Oid	   Id;
				IoHash Hash;
			};

			struct FileMapping
			{
				Oid			Id;
				IoHash		Hash;		 // This is either zero or a cid
				std::string ServerPath;	 // If Hash is valid then this should be empty
				std::string ClientPath;
			};

			std::vector<Mapping>	 Chunks;
			std::vector<Mapping>	 Meta;
			std::vector<FileMapping> Files;
		};

		OplogEntryMapping GetMapping(CbObjectView Core);

		/** Update tracking metadata for a new oplog entry
		 *
		 * This is used during replay (and gets called as part of new op append)
		 *
		 * Returns the oplog LSN assigned to the new entry, or kInvalidOp if the entry is rejected
		 */
		uint32_t RegisterOplogEntry(RwLock::ExclusiveLockScope& OplogLock, const OplogEntryMapping& OpMapping, const OplogEntry& OpEntry);

		void AddFileMapping(const RwLock::ExclusiveLockScope& OplogLock,
							const Oid&						  FileId,
							const IoHash&					  Hash,
							std::string_view				  ServerPath,
							std::string_view				  ClientPath);
		void AddChunkMapping(const RwLock::ExclusiveLockScope& OplogLock, const Oid& ChunkId, const IoHash& Hash);
		void AddMetaMapping(const RwLock::ExclusiveLockScope& OplogLock, const Oid& ChunkId, const IoHash& Hash);

		friend class ProjectStoreOplogReferenceChecker;
		friend class ProjectStoreReferenceChecker;
	};

	struct Project : public RefCounted
	{
		std::string			  Identifier;
		std::filesystem::path RootDir;
		std::filesystem::path EngineRootDir;
		std::filesystem::path ProjectRootDir;
		std::filesystem::path ProjectFilePath;

		Oplog*					 NewOplog(std::string_view OplogId, const std::filesystem::path& MarkerPath);
		Oplog*					 OpenOplog(std::string_view OplogId);
		void					 DeleteOplog(std::string_view OplogId);
		std::filesystem::path	 RemoveOplog(std::string_view OplogId);
		void					 IterateOplogs(std::function<void(const RwLock::SharedLockScope&, const Oplog&)>&& Fn) const;
		void					 IterateOplogs(std::function<void(const RwLock::SharedLockScope&, Oplog&)>&& Fn);
		std::vector<std::string> ScanForOplogs() const;
		bool					 IsExpired(const RwLock::SharedLockScope&, const GcClock::TimePoint ExpireTime);
		bool			   IsExpired(const RwLock::SharedLockScope&, const GcClock::TimePoint ExpireTime, const ProjectStore::Oplog& Oplog);
		bool			   IsExpired(const GcClock::TimePoint ExpireTime, const ProjectStore::Oplog& Oplog);
		void			   TouchProject() const;
		void			   TouchOplog(std::string_view Oplog) const;
		GcClock::TimePoint LastOplogAccessTime(std::string_view Oplog) const;

		Project(ProjectStore* PrjStore, CidStore& Store, std::filesystem::path BasePath);
		virtual ~Project();

		void					  Read();
		void					  Write();
		[[nodiscard]] static bool Exists(const std::filesystem::path& BasePath);
		void					  Flush();
		void					  ScrubStorage(ScrubContext& Ctx);
		LoggerRef				  Log();
		void					  GatherReferences(GcContext& GcCtx);
		static uint64_t			  TotalSize(const std::filesystem::path& BasePath);
		uint64_t				  TotalSize() const;
		bool					  PrepareForDelete(std::filesystem::path& OutDeletePath);

		void					 EnableUpdateCapture();
		void					 DisableUpdateCapture();
		std::vector<std::string> GetCapturedOplogs();

		std::vector<RwLock::SharedLockScope> GetGcReferencerLocks();

	private:
		ProjectStore*									   m_ProjectStore;
		CidStore&										   m_CidStore;
		mutable RwLock									   m_ProjectLock;
		std::map<std::string, std::unique_ptr<Oplog>>	   m_Oplogs;
		std::vector<std::unique_ptr<Oplog>>				   m_DeletedOplogs;
		std::filesystem::path							   m_OplogStoragePath;
		mutable tsl::robin_map<std::string, GcClock::Tick> m_LastAccessTimes;
		mutable RwLock									   m_UpdateCaptureLock;
		uint32_t										   m_UpdateCaptureRefCounter = 0;
		std::unique_ptr<std::vector<std::string>>		   m_CapturedOplogs;

		std::filesystem::path BasePathForOplog(std::string_view OplogId);
		bool				  IsExpired(const RwLock::SharedLockScope&,
										const std::string&			 EntryName,
										const std::filesystem::path& MarkerPath,
										const GcClock::TimePoint	 ExpireTime);
		void				  WriteAccessTimes();
		void				  ReadAccessTimes();

		friend class ProjectStoreOplogReferenceChecker;
		friend class ProjectStoreReferenceChecker;
	};

	//	Oplog* OpenProjectOplog(std::string_view ProjectId, std::string_view OplogId);

	Ref<Project> OpenProject(std::string_view ProjectId);
	Ref<Project> NewProject(const std::filesystem::path& BasePath,
							std::string_view			 ProjectId,
							const std::filesystem::path& RootDir,
							const std::filesystem::path& EngineRootDir,
							const std::filesystem::path& ProjectRootDir,
							const std::filesystem::path& ProjectFilePath);
	bool		 UpdateProject(std::string_view				ProjectId,
							   const std::filesystem::path& RootDir,
							   const std::filesystem::path& EngineRootDir,
							   const std::filesystem::path& ProjectRootDir,
							   const std::filesystem::path& ProjectFilePath);
	bool		 RemoveProject(std::string_view ProjectId, std::filesystem::path& OutDeletePath);
	bool		 DeleteProject(std::string_view ProjectId);
	bool		 Exists(std::string_view ProjectId);
	void		 Flush();
	void		 DiscoverProjects();
	void		 IterateProjects(std::function<void(Project& Prj)>&& Fn);

	LoggerRef					 Log() { return m_Log; }
	const std::filesystem::path& BasePath() const { return m_ProjectBasePath; }

	virtual void		  GatherReferences(GcContext& GcCtx) override;
	virtual void		  ScrubStorage(ScrubContext& Ctx) override;
	virtual void		  CollectGarbage(GcContext& GcCtx) override;
	virtual GcStorageSize StorageSize() const override;

	virtual std::string						 GetGcName(GcCtx& Ctx) override;
	virtual GcStoreCompactor*				 RemoveExpiredData(GcCtx& Ctx, GcStats& Stats) override;
	virtual std::vector<GcReferenceChecker*> CreateReferenceCheckers(GcCtx& Ctx) override;

	virtual std::vector<RwLock::SharedLockScope> LockState(GcCtx& Ctx) override;

	CbArray									 GetProjectsList();
	std::pair<HttpResponseCode, std::string> GetProjectFiles(const std::string_view					ProjectId,
															 const std::string_view					OplogId,
															 const std::unordered_set<std::string>& WantedFieldNames,
															 CbObject&								OutPayload);
	std::pair<HttpResponseCode, std::string> GetProjectChunkInfos(const std::string_view				 ProjectId,
																  const std::string_view				 OplogId,
																  const std::unordered_set<std::string>& WantedFieldNames,
																  CbObject&								 OutPayload);
	std::pair<HttpResponseCode, std::string> GetChunkInfo(const std::string_view ProjectId,
														  const std::string_view OplogId,
														  const std::string_view ChunkId,
														  CbObject&				 OutPayload);
	std::pair<HttpResponseCode, std::string> GetChunkRange(const std::string_view ProjectId,
														   const std::string_view OplogId,
														   const Oid			  ChunkId,
														   uint64_t				  Offset,
														   uint64_t				  Size,
														   ZenContentType		  AcceptType,
														   CompositeBuffer&		  OutChunk,
														   ZenContentType&		  OutContentType);
	std::pair<HttpResponseCode, std::string> GetChunkRange(const std::string_view ProjectId,
														   const std::string_view OplogId,
														   const std::string_view ChunkId,
														   uint64_t				  Offset,
														   uint64_t				  Size,
														   ZenContentType		  AcceptType,
														   CompositeBuffer&		  OutChunk,
														   ZenContentType&		  OutContentType);
	std::pair<HttpResponseCode, std::string> GetChunk(const std::string_view ProjectId,
													  const std::string_view OplogId,
													  const std::string_view Cid,
													  ZenContentType		 AcceptType,
													  IoBuffer&				 OutChunk);

	std::pair<HttpResponseCode, std::string> PutChunk(const std::string_view ProjectId,
													  const std::string_view OplogId,
													  const std::string_view Cid,
													  ZenContentType		 ContentType,
													  IoBuffer&&			 Chunk);

	std::pair<HttpResponseCode, std::string> WriteOplog(const std::string_view ProjectId,
														const std::string_view OplogId,
														IoBuffer&&			   Payload,
														CbObject&			   OutResponse);

	std::pair<HttpResponseCode, std::string> ReadOplog(const std::string_view				 ProjectId,
													   const std::string_view				 OplogId,
													   const HttpServerRequest::QueryParams& Params,
													   CbObject&							 OutResponse);

	bool Rpc(HttpServerRequest&		HttpReq,
			 const std::string_view ProjectId,
			 const std::string_view OplogId,
			 IoBuffer&&				Payload,
			 AuthMgr&				AuthManager);

	std::pair<HttpResponseCode, std::string> Export(Ref<ProjectStore::Project> Project,
													ProjectStore::Oplog&	   Oplog,
													CbObjectView&&			   Params,
													AuthMgr&				   AuthManager);

	std::pair<HttpResponseCode, std::string> Import(ProjectStore::Project& Project,
													ProjectStore::Oplog&   Oplog,
													CbObjectView&&		   Params,
													AuthMgr&			   AuthManager);

	bool AreDiskWritesAllowed() const;

	void					 EnableUpdateCapture();
	void					 DisableUpdateCapture();
	std::vector<std::string> GetCapturedProjects();

private:
	LoggerRef								  m_Log;
	GcManager&								  m_Gc;
	CidStore&								  m_CidStore;
	JobQueue&								  m_JobQueue;
	std::filesystem::path					  m_ProjectBasePath;
	mutable RwLock							  m_ProjectsLock;
	std::map<std::string, Ref<Project>>		  m_Projects;
	const DiskWriteBlocker*					  m_DiskWriteBlocker = nullptr;
	mutable RwLock							  m_UpdateCaptureLock;
	uint32_t								  m_UpdateCaptureRefCounter = 0;
	std::unique_ptr<std::vector<std::string>> m_CapturedProjects;

	std::filesystem::path BasePathForProject(std::string_view ProjectId);

	friend class ProjectStoreGcStoreCompactor;
	friend class ProjectStoreOplogReferenceChecker;
	friend class ProjectStoreReferenceChecker;
};

void prj_forcelink();

}  // namespace zen
